package middleware

import (
	"SmartLibrary/pkg/log"
	"github.com/gin-gonic/gin"
	"github.com/juju/ratelimit"
	"time"
)

// RateLimit 限流中间件
func RateLimit(logger *log.Logger, limit int, refillInterval time.Duration) gin.HandlerFunc {
	bucket := ratelimit.NewBucket(refillInterval, int64(limit))

	return func(ctx *gin.Context) {
		// 尝试从桶中获取令牌
		if bucket.TakeAvailable(1) < 0 {
			// 如果没有可用的令牌，则记录日志并返回
			logger.Error("too many requests")
			ctx.Abort()
			return
		}
		ctx.Next()
	}
}
